import numpy as np
import torch
import torch.nn as nn
from torch.distributions import Categorical, MultivariateNormal

################################## set device ##################################
print("============================================================================================")
# set device to cpu or cuda
device = torch.device('cpu')
if(torch.cuda.is_available()): 
    device = torch.device('cuda:0') 
    torch.cuda.empty_cache()
    print("Device set to : " + str(torch.cuda.get_device_name(device)))
else:
    print("Device set to : cpu")
print("============================================================================================")


################################## PPO Policy ##################################
class RolloutBuffer:
    def __init__(self):
        self.actions = []
        self.states = []
        self.logprobs = []
        self.rewards = []
        self.state_values = []
        self.is_terminals = []
    
    def clear(self):
        del self.actions[:]
        del self.states[:]
        del self.logprobs[:]
        del self.rewards[:]
        del self.state_values[:]
        del self.is_terminals[:]


class ActorCritic(nn.Module):
    def __init__(self, state_dim:int, action_dim:int, has_continuous_action_space:bool, action_std_init:float) -> None:
        """
        ActorCritic 网络模型

        Args:
            state_dim (int): 状态空间的数量
            action_dim (int): 行动空间的数量
            has_continuous_action_space (bool): 是否有连续动作空间
            action_std_init (float): 如果是连续动作空间，设置初始的动作标准差
        """
        super(ActorCritic, self).__init__()
        # 这里是演员评论家网络模型，继承于 nn.Module
        self.has_continuous_action_space = has_continuous_action_space    # self.has_continuous_action_space: bool -> 是否是连续动作空间的策略
        
        """
        如果是连续动作空间
        就生成一个 action_var，size: [action_dim]
        连续动作空间使用的是 MultivariateNormal 多元正态分布，需要一个协方差矩阵，这里用 action_std_init * action_std_init 填充
        输出的 Tensor size 是 [action_dim]
        """
        if has_continuous_action_space:
            self.action_dim = action_dim
            self.action_var = torch.full((action_dim,), action_std_init * action_std_init).to(device)    # 这里是将动作空间用 action_std_init*action_std_init 填充起来
        # actor
        if has_continuous_action_space :
            self.actor = nn.Sequential(
                            nn.Linear(state_dim, 64),
                            nn.Tanh(),
                            nn.Linear(64, 64),
                            nn.Tanh(),
                            nn.Linear(64, action_dim),
                            nn.Tanh()
                        )    # 这里是全连接网络，作为主要的 Actor 网络，nn.Tanh() 函数是双正切激活函数
        else:
            self.actor = nn.Sequential(
                            nn.Linear(state_dim, 64),
                            nn.Tanh(),
                            nn.Linear(64, 64),
                            nn.Tanh(),
                            nn.Linear(64, action_dim),
                            nn.Softmax(dim=-1)
                        )
        # critic
        self.critic = nn.Sequential(
                        nn.Linear(state_dim, 64),
                        nn.Tanh(),
                        nn.Linear(64, 64),
                        nn.Tanh(),
                        nn.Linear(64, 1)
                    )       # 这里是 critic 评论家网络，为什么最后输出的维度是 1? 因为我们只需要输出一个值，即评估值。
        
    def set_action_std(self, new_action_std: float) -> None:
        """
        设置动作空间的标准差
        只有对连续动作空间有效

        Args:
            new_action_std (float): 新的动作控价标准差
        """
        if self.has_continuous_action_space:
            self.action_var = torch.full((self.action_dim,), new_action_std * new_action_std).to(device)
        else:
            print("--------------------------------------------------------------------------------------------")
            print("WARNING : Calling ActorCritic::set_action_std() on discrete action space policy")
            print("--------------------------------------------------------------------------------------------")

    def forward(self) -> None:
        raise NotImplementedError
    
    def act(self, state):
        """
        策略网络（通常是 Actor 网络）来执行

        Args:
            state (torch.Tensor): 状态空间的输入

        Returns:
            action_distribution(torch.Tensor): 根据当前状态 $$s_t$$ 返回动作的概率分布
            log_probability_of_actions(torch.Tensor)：返回在给定状态下选择某个动作的对数概率，对数概率用于计算重要性采样比率和 PPO 的损失函数
            value_of_state(torch.Tensor): 如果包含 critic 网络，会返回状态值函数 $$V(s_t)$$，用于计算优势函数
        """
        # 在 PPO 算法的 select_action() 函数中，先让 
        if self.has_continuous_action_space:                         # 如果是连续动作空间的话，首先直接推理拿到 action_mean，
            action_mean = self.actor(state)                          # 在 ActorCritic 网络中，先让 actor 来 基于 state 推理 action_mean 也就是 action 的期望值
            cov_mat = torch.diag(self.action_var).unsqueeze(dim=0)   # torch.diag() 函数是根据一维张量创建对角矩阵，然后再使用 unsqueeze() 函数增加维度，增加后的维度应该是 (1, 6, 6)
            dist = MultivariateNormal(action_mean, cov_mat)          # 多元正态分布
        else:
            action_probs = self.actor(state)
            dist = Categorical(action_probs)

        action = dist.sample()      # action: torch.Tensor -> 使用多元正态分布进行采样，得到 action 的值
        action2 = dist.sample()
        action_logprob = dist.log_prob(action)        # log_prob: 生成对数概率，ln 概率的好处
        state_val = self.critic(state)                # state_val: torch.Tensor -> 评论家网络的评估值

        return action.detach(), action_logprob.detach(), state_val.detach()     # detach 函数的功能是避免梯度计算
    
    def evaluate(self, state, action):
        """
        在 PPO (Proximal Policy Optimization) 中，Actor-Critic 网络的 evaluate 函数用于
        评估当前策略在给定状态下的动作分布和状态价值函数。这个函数通常在策略更新阶段被调用，
        用于计算新策略的动作对数概率和状态值函数，从而进一步计算重要性采样比率、优势函数和损失
        函数

        Args:
            state (torch.Tensor): size: [批次的数量，状态空间的维度]; 含义: 批次的状态空间输入
            action (torch.Tensor): size: [批次的数量，动作空间的维度]; 含义: 批次的动作空间输入

        Returns:
            action_logprobs (torch.Tensor): size: [批次的数量]; 含义: 用 actor 网络计算出的动作的对数概率
            state_values (torch.Tensor): size: [批次的数量]; 含义: 用 critic 网络计算出的状态价值函数 $$V(s_t)$$
            dist_entropy (torch.tensor): size: [批次的数量]; 含义: 用 actor 网络计算出动作分布的熵
        """
        # state 和 action 都是一个批次的 tensor，state 的 size 为 [4000, 27]，action 的 size 为 [4000, 6]
        if self.has_continuous_action_space:   # 如果是连续动作空间
            action_mean = self.actor(state)    # 1. 计算动作分布 action_mean: torch.Tensor -> 基于 state 推理得到 action 的期望值，这里的 state 是一个批次的 state，size 为[4000, 27]
            
            action_var = self.action_var.expand_as(action_mean)    # 此时 action_var 的 size 为 [4000, 6]
            cov_mat = torch.diag_embed(action_var).to(device)      # diag_embed() 函数是将输入张量嵌入到一个更高维张量的对角线上，从而创建出对角张量，此时 cov_mat 的 size [4000, 6, 6]
            dist = MultivariateNormal(action_mean, cov_mat)        # 多元正态分布
            
            # For Single Action Environments.
            if self.action_dim == 1:
                action = action.reshape(-1, self.action_dim)
        else:
            action_probs = self.actor(state)
            dist = Categorical(action_probs)
        action_logprobs = dist.log_prob(action)                    # action_logprobs: torch.Tensor -> 计算动作的对数概率
        dist_entropy = dist.entropy()                              # dist_entropy: torch.Tensor -> 动作分布的熵，用于衡量动作分布的随机性
        state_values = self.critic(state)                          # state_values: torch.Tensor -> 评论家网络的评估值，也就是状态价值函数 V(s_t)
        
        return action_logprobs, state_values, dist_entropy


class PPO:
    def __init__(
            self, 
            state_dim:int, 
            action_dim:int, 
            lr_actor:float, 
            lr_critic:float, 
            gamma:float, 
            K_epochs:int, 
            eps_clip:float, 
            has_continuous_action_space:bool, 
            action_std_init=0.6) -> None:
        """
        PPO Actor Critic 的最小实现

        Args:
            state_dim (int): 状态空间的维度
            action_dim (int): 动作空间的维度
            lr_actor (float): Actor 网络的学习率
            lr_critic (float): Critic 网络的学习率
            gamma (float): 折扣因子
            K_epochs (int): 对数据收集器中的数据进行优化的次数
            eps_clip (float): reward 的折扣系数
            has_continuous_action_space (bool): 是否是连续动作空间
            action_std_init (float, optional): 如果是连续动作空间，设置初始的动作标准差. Defaults to 0.6.
        """
        # 这部分是 PPO 网络的最小实现
        self.has_continuous_action_space = has_continuous_action_space                   # self.has_continuous_action_space: bool -> 是否是连续动作空间的策略

        if has_continuous_action_space:                                                  # 如果是连续动作空间
            self.action_std = action_std_init                                            # self.action_std: float -> 就设置初始的动作空间输出概率分布为 xx

        self.gamma = gamma                                                               # self.gamma: float -> 奖励的折扣系数
        self.eps_clip = eps_clip                                                         # self.eps_clip: float -> PPO 算法中的 epsilon 值，用于限制 policy 更新的幅度
        self.K_epochs = K_epochs                                                         # self.K_epochs: int -> K_epochs 是指在优化的时候同一批数据优化多少次
        
        self.buffer = RolloutBuffer()                                                    # self.buffer: RolloutBuffer -> 用于存储策略的训练数据

        self.policy = ActorCritic(state_dim, action_dim, has_continuous_action_space, action_std_init).to(device)     # 这个是 PPO 的主要的 policy 网络
        self.optimizer = torch.optim.Adam([
                        {'params': self.policy.actor.parameters(), 'lr': lr_actor},
                        {'params': self.policy.critic.parameters(), 'lr': lr_critic}
                    ])         # 这里是 PPO 算法的优化器，使用 Adam 优化器，优化器的全称是 Adaptive Moment Estimation，能够自适应的为不同的参数调整学习率。

        self.policy_old = ActorCritic(state_dim, action_dim, has_continuous_action_space, action_std_init).to(device)     # 这个是 PPO 的 old policy 网络，用于存储之前的 policy 网络参数，用于计算策略的更新
        self.policy_old.load_state_dict(self.policy.state_dict())                                                         # 这个的意思是让 policy_old 的参数和 policy 的参数相同
        
        self.MseLoss = nn.MSELoss()    # nn.MSELoss() 是均方差损失，Mean Squared Error Loss，用于评估 critic 网络的输出值和实际奖励的差距。

    def set_action_std(self, new_action_std: float) -> None:
        """
        设置动作空间的标准差

        Args:
            new_action_std (float): 新的动作空间的标准差
        """
        if self.has_continuous_action_space:
            self.action_std = new_action_std
            self.policy.set_action_std(new_action_std)
            self.policy_old.set_action_std(new_action_std)
        else:
            print("--------------------------------------------------------------------------------------------")
            print("WARNING : Calling PPO::set_action_std() on discrete action space policy")
            print("--------------------------------------------------------------------------------------------")

    def decay_action_std(self, action_std_decay_rate:float, min_action_std:float) -> None:
        """
        减少动作空间输出的标准差

        Args:
            action_std_decay_rate (float): 动作空间标准差的衰减率的折扣系数
            min_action_std (float): 最小的动作空间标准差
        """
        print("--------------------------------------------------------------------------------------------")
        if self.has_continuous_action_space:
            self.action_std = self.action_std - action_std_decay_rate
            self.action_std = round(self.action_std, 4)
            if (self.action_std <= min_action_std):
                self.action_std = min_action_std
                print("setting actor output action_std to min_action_std : ", self.action_std)
            else:
                print("setting actor output action_std to : ", self.action_std)
            self.set_action_std(self.action_std)

        else:
            print("WARNING : Calling PPO::decay_action_std() on discrete action space policy")
        print("--------------------------------------------------------------------------------------------")

    def select_action(self, state:np.ndarray) -> np.ndarray:
        """
        根据输入的状态空间选择 state 对应的 action space

        Args:
            state (np.ndarray): 状态空间的输入

        Returns:
            np.ndarray: 动作空间输出
        """
        # 根据输入的状态空间的数据 state，选择对应的 action space
        if self.has_continuous_action_space:
            with torch.no_grad():            # 在 select_action 的时候，是不需要计算梯度的，所以这里使用 torch.no_grad() 装饰器
                state = torch.FloatTensor(state).to(device)     # state:torch.Tensor -> 这里将 np.ndarray 转换为 torch.Tensor
                action, action_logprob, state_val = self.policy_old.act(state)    # 1. 先让 policy_old.act 基于 state 来 act action

            self.buffer.states.append(state)                    # 在经验回放缓冲区保存 state，本时刻的状态
            self.buffer.actions.append(action)                  # 在经验回放缓冲区保存 action，本时刻的动作概率
            self.buffer.logprobs.append(action_logprob)         # 本时刻 action_logprob，本时刻的动作对数概率
            self.buffer.state_values.append(state_val)          # 本时刻的状态价值函数

            return action.detach().cpu().numpy().flatten()      # 但是只返回 action 的值用 fltten 是展平为 size 为 [1]
        else:
            with torch.no_grad():
                state = torch.FloatTensor(state).to(device)
                action, action_logprob, state_val = self.policy_old.act(state)
            
            self.buffer.states.append(state)
            self.buffer.actions.append(action)
            self.buffer.logprobs.append(action_logprob)
            self.buffer.state_values.append(state_val)

            return action.item()

    def update(self) -> None:
        """
        更新网络模型的参数
        """
        # Monte Carlo estimate of returns  蒙特卡洛估计
        rewards = []                                            # rewards: list -> 
        discounted_reward = 0                                   # discounted_reward: float -> 折扣后的 reward
        for reward, is_terminal in zip(reversed(self.buffer.rewards), reversed(self.buffer.is_terminals)):    # 对于倒着的奖励和是否终止的列表，逐个取出
            if is_terminal:
                discounted_reward = 0
            discounted_reward = reward + (self.gamma * discounted_reward)       # 如果倒着数的 is_terminal 是 True，那么 discounted_reward 就置为 0，否则加上 gamma 乘以之前的 discounted_reward
            rewards.insert(0, discounted_reward)                # 
            
        # Normalizing the rewards
        rewards = torch.tensor(rewards, dtype=torch.float32).to(device)    # rewards: torch.Tensor -> 转换为 torch.Tensor，此时 rewords 的 size 为 [4000]
        rewards = (rewards - rewards.mean()) / (rewards.std() + 1e-7)      # 标准化奖励

        # convert list to tensor，将经验回放缓冲区中的数据全部转换为 Tensor
        old_states = torch.squeeze(torch.stack(self.buffer.states, dim=0)).detach().to(device)              # old_states:torch.Tensor -> 正向顺序排布的 action_net 生成的 states
        old_actions = torch.squeeze(torch.stack(self.buffer.actions, dim=0)).detach().to(device)            # old_actions:torch.Tensor -> 正向顺序排布的 action_net 生成的 actions
        old_logprobs = torch.squeeze(torch.stack(self.buffer.logprobs, dim=0)).detach().to(device)          # old_logprobs:torch.Tensor -> 正向顺序排布的 action_net 生成的动作对数概率
        old_state_values = torch.squeeze(torch.stack(self.buffer.state_values, dim=0)).detach().to(device)  # old_state_values:torch.Tensor -> 正向顺序排布的 critic_net 生成的状态价值函数

        # calculate advantages
        advantages = rewards.detach() - old_state_values.detach()    # advantage: torch.Tensor -> 计算优势，优势是“状态-动作值函数”计算出来的未来累计奖励的期望值 - 状态价值函数（就是之前 Critic 网络根据过去的 state 算出来的）

        # Optimize policy for K epochs 这个 K_epochs 是指对同一批数据优化多少次
        for _ in range(self.K_epochs):

            # Evaluating old actions and values
            logprobs, state_values, dist_entropy = self.policy.evaluate(old_states, old_actions)

            # match state_values tensor dimensions with rewards tensor
            state_values = torch.squeeze(state_values)    # 将 状态价值 和 rewards 的维度处理成相同的，squeeze 函数是去除 Tensor 中维度维 1 的维度
            
            # Finding the ratio (pi_theta / pi_theta__old)
            """
            ratios 是重要性比率，该比率的定义是新策略的动作概率与旧概率的动作概率之比
            logprobs 是新策略的动作概率
            old_logprobs 是旧策略的动作概率
            这个比率用于计算重要性采样（importance sampling）以及限制策略更新的幅度
            比率的计算过程是：$$log 这两个数求差 $$
            1. old_log：在之前的采样阶段，从回放缓冲区中取出的动作对数概率
            2. logprobs: 使用当前的 policy 计算出来的相同动作的对数概率
            """
            ratios = torch.exp(logprobs - old_logprobs.detach())    # ratios: torch.Tensor(size: [4000]) -> 重要性比率

            # Finding Surrogate Loss  计算替代损失
            """
            替代损失的作用：
            1. 优化策略：通过最大化替代损失，可以改进策略的性能。
            2. 限制策略更新幅度：通过裁剪机制（clipping mechanism）,可以避免新策略偏离旧策略太远，从而保证训练的稳定性
            
            计算替代损失的过程
            1. 计算重要性采样比率，这个前面已经计算出来了，就是 ratios
            2. 计算未裁剪的目标值：surr1 = ratios * advantages
            3. 计算裁剪后的目标值：surr2 = 限制 ratios 在 1-self.eps_clip 和 1+self.eps_clip 之间，然后乘以 优势函数
            4. 最终的替代损失 -torch.min(surr1, surr2) 
            """
            surr1 = ratios * advantages
            surr2 = torch.clamp(ratios, 1-self.eps_clip, 1+self.eps_clip) * advantages    # surr2:替代2 torch.clamp() 函数是将 ratios 的值限制在 1-self.eps_clip 和 1+self.eps_clip 之间，用于限制策略更新的梯度

            # final loss of clipped objective PPO
            """
            PPO 的最终裁剪损失由 3 个部分组成
            1. -torch.min(surr1, surr2)：裁剪目标函数，同时限制策略更新的幅度，计算方法如上
            2. 0.5 * self.MseLoss(state_values, rewards)：价值函数损失，使用均方差损失计算 状态价值函数 和 目标值函数，前面再乘以一个系数
            3. - 0.01 * dist_entropy：熵正则化损失，孤立策略保持一定的随机性
            #! TODO 这里的 价值函数损失和熵正则化损失的系数正负号是否正确？
            """
            loss = -torch.min(surr1, surr2) + 0.5 * self.MseLoss(state_values, rewards) - 0.01 * dist_entropy
            
            # take gradient step
            """
            最后梯度更新的步骤
            1. 首先清空之前的梯度
            2. 计算损失函数的平均值，反向传播梯度
            为什么要先求平均：
                1. 前面的损失是由一个批次的损失出来的，有多个样本，而不是单个样本，通过对这个损失张量
                求平均值，可以将批量中所有样本的损失合并为一个标量值，好处是可以转化为一个统一的优化
                问题。
                2. 与优化器的兼容性
            3. 优化器.step() 函数
                这个是根据计算得到的梯度，使用对应的优化算法来更新模型参数，如果不调用，这些梯度只会
                保存在参数的 .grad 属性中。
            """
            self.optimizer.zero_grad()     # 这里的 optimizer.zero_grad() step 是清空之前的梯度，防止梯度累加，梯度累加是 pytorch 中的一个机制
            loss_mean = loss.mean()        #! 为什么这里要先求平均再反向传播
            loss_mean.backward()
            # loss.mean().backward()
            self.optimizer.step()
            
        # Copy new weights into old policy
        self.policy_old.load_state_dict(self.policy.state_dict())    # 将新策略的权重复制到旧策略中

        # clear buffer
        self.buffer.clear()    # 清空 Buffer
    
    def save(self, checkpoint_path:str) -> None:
        """
        保存模型参数到文件中

        Args:
            checkpoint_path (str): 保存到哪里
        """
        torch.save(self.policy_old.state_dict(), checkpoint_path)
   
    def load(self, checkpoint_path:str):
        """
        从文件中加载模型参数到 policy_old 以及 policy 网络中

        Args:
            checkpoint_path (str): checkpoint 文件的路径
        """
        self.policy_old.load_state_dict(torch.load(checkpoint_path, map_location=lambda storage, loc: storage))
        self.policy.load_state_dict(torch.load(checkpoint_path, map_location=lambda storage, loc: storage))
        
        
       


